import re  # 需导入re模块

def clean_scenic_info(info):
    """景点简介文本清洗函数：保留有效语义，去除噪声"""
    if not info or info.strip() == "":
        return ""
    
    # 1. 去除特殊符号、乱码（保留中文、英文、数字、常见标点）
    cleaned = re.sub(
        r'[^\u4e00-\u9fa5a-zA-Z0-9\s，。！？；：""''（）【】(){}、]', 
        '', 
        info
    )
    
    # 2. 合并冗余空格，去除首尾空格
    cleaned = re.sub(r'\s+', ' ', cleaned).strip()
    
    # 3. 去除重复片段（简单版：如连续重复的3个以上字符/词语）
    # （复杂场景可使用difflib或LCS算法，此处为轻量版）
    cleaned = re.sub(r'(\w{2,})\1{2,}', r'\1', cleaned)  # 如“夫子庙夫子庙夫子庙”→“夫子庙”
    
    # 4. 处理过长文本（BERT默认max_length=512，此处按“。”拆分取前3句）
    if len(cleaned) > 500:  # 预留12个字符给BERT的特殊符号（[CLS]、[SEP]）
        sentences = re.split(r'[。！？；]', cleaned)
        cleaned = "。".join(sentences[:3]) + "。"  # 保留前3句，补全句号
    
    return cleaned

# 在generate_and_store_embeddings函数中调用清洗函数
def generate_and_store_embeddings():
    print("正在加载BERT模型...")
    scorer = DualRelationScorer(model_dir)
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()
    
    try:
        cursor.execute("SELECT id, info FROM scenic_spots WHERE embedding IS NULL")
        spots = cursor.fetchall()
        total = len(spots)
        print(f"找到 {total} 个景点需要生成embedding")
        
        for i, (spot_id, info) in enumerate(spots, 1):
            # --- 新增：文本清洗 ---
            cleaned_info = clean_scenic_info(info)
            if not cleaned_info:  # 清洗后为空，跳过
                print(f"景点 {spot_id} 清洗后简介为空，跳过")
                continue
            
            try:
                # 使用清洗后的文本生成embedding
                embedding = scorer.extract_single_text_embedding(cleaned_info)
                embedding_blob = embedding.tobytes()
                
                cursor.execute(
                    "UPDATE scenic_spots SET embedding = ? WHERE id = ?",
                    (embedding_blob, spot_id)
                )
                
                if i % 10 == 0:
                    conn.commit()
                    print(f"已处理 {i}/{total} 个景点")
            except Exception as e:
                print(f"处理景点 {spot_id} 时出错: {str(e)}")
                continue
        
        conn.commit()
        print(f"所有景点embedding生成完成，共处理 {total} 个景点")
    except Exception as e:
        print(f"生成embedding时出错: {str(e)}")
        conn.rollback()
    finally:
        conn.close()